package limit

import (
	"github.com/gin-gonic/gin"
	"hobby-mate/common"
	"net/http"
	"sync"
)

var tokenBucket *TokenBucket
var mutex sync.Mutex

// ApiLimit 接口限流
// 在global后执行本中间件
func ApiLimit(routerServer *gin.Context) {
	mutex.Lock()
	defer mutex.Unlock()
	if tokenBucket == nil {
		tokenBucket = SimpleApiLimit(&(TokenBucketConfig{
			Capacity:   10,
			RefillRate: 2,
			OkCallBack: func() { routerServer.Next() },
			FailCallBack: func() {
				routerServer.AbortWithStatusJSON(http.StatusTooManyRequests, common.TooManyRequestErrorJSON)
			},
		}))
	} else {
		tokenBucket.TakeToken()
	}
}
